#!/usr/bin/env bash
#filename: remk_image.sh

#set -x

SCRIPT_NAME=$(basename "$0")
SCRIPT_DIR=$(cd "$(dirname "$0")";pwd)
CUR_DIR=$(pwd)

Usage1() {
	echo "Usage: ${SCRIPT_NAME} [hexzci:r:s:k:o:]"
	echo "Options:"
	echo "    -h --- This help text"
	echo "    -e --- Extract rootfs from image"
	echo "    -x --- Unpack squashfs"
	echo "    -z --- Pack squashfs"
	echo "    -c --- Create new image from rootfs and vminux"
	echo "    -i <image_path_name> --- image file path and name"
	echo "    -r <rootfs_path_name> --- rootfs file path and name"
	echo "    -s <squashfs_path> --- Squashfs path"
	echo "    -k <vminux_path_name> --- vminux file path and name"
	echo "    -o <output_path> --- Output file path(rootfs or image)"
	echo "Examples:"
	echo "    ${SCRIPT_NAME} -e -i ./acu100e/acu100e-image.bin -o ./acu100e"
	echo "    ${SCRIPT_NAME} -x -r ./acu100e/acu100e-rootfs.bin -o ./acu100e"
	echo "    ${SCRIPT_NAME} -z -s ./acu100e/squashfs-root -o ./acu100e"
	echo "    ${SCRIPT_NAME} -c -k ./acu100e/acu100e-vmlinux.bin -r ./acu100e/acu100e-new-rootfs.bin -o ./acu100e"
	return 0
}

Usage2() {
	echo "用法: ${SCRIPT_NAME} [hexzci:r:s:k:o:]"
	echo "选项:"
	echo "    -h --- 显示帮助信息"
	echo "    -e --- 从镜像中提取rootfs"
	echo "    -x --- 解包squashfs"
	echo "    -z --- 打包squashfs"
	echo "    -c --- 从rootfs和vminux创建新image"
	echo "    -i <image_path_name> --- image文件路径和文件名"
	echo "    -r <rootfs_path_name> --- rootfs文件路径和文件名"
	echo "    -s <squashfs_path> --- squashfs路径"
	echo "    -k <vminux_path_name> --- vminux文件路径和文件名"
	echo "    -o <output_path> --- 输出文件路径(rootfs或image)"
	echo "示例:"
	echo "    ${SCRIPT_NAME} -e -i ./acu100e/acu100e-image.bin -o ./acu100e"
	echo "    ${SCRIPT_NAME} -x -r ./acu100e/acu100e-rootfs.bin -o ./acu100e"
	echo "    ${SCRIPT_NAME} -z -s ./acu100e/squashfs-root -o ./acu100e"
	echo "    ${SCRIPT_NAME} -c -k ./acu100e/acu100e-vmlinux.bin -r ./acu100e/acu100e-new-rootfs.bin -o ./acu100e"
	return 0
}

Usage() {
	local iszh=$(env | grep "LANG=" | grep "zh")
	if [ -z "${iszh}"  ]; then
		Usage1
	else
		Usage2
	fi
}

work_mode=''
output_path=''
output_vmlinux=''
output_rootfs=''
output_squashfs=''
output_image=''

input_image=''
input_vmlinux=''
input_rootfs=''
input_squashfs=''

if [ $# -lt 1 ]; then
	Usage
	exit 1
fi

while getopts "hexzci:r:s:k:o:" opt ; do
	case "$opt" in
	'h')
		Usage
		exit 0
		;;
	'e')
		work_mode='extract'
		output_vmlinux='output_vmlinux.bin'
		output_rootfs='output_rootfs.bin'
		;;
	'x')
		work_mode='unpack'
		output_squashfs='squashfs-root'
		;;
	'z')
		work_mode='pack'
		output_rootfs='output_new_rootfs.bin'
		;;
	'c')
		work_mode='create'
		output_image='output_new_image.bin'
		;;
	'i')
		input_image="${OPTARG}"
		;;
	'r')
		input_rootfs="${OPTARG}"
		;;
	's')
		input_squashfs="${OPTARG}"
		;;
	'k')
		input_vmlinux="${OPTARG}"
		;;
	'o')
		output_path="${OPTARG}"
		;;
	*)
		Usage
		exit 1
		;;
	esac
done

die()
{
	echo "Die: $@"
	exit 15
}

extract_rootfs()
{
	if [ $# -ne 4 ]; then
		echo "Error: parameter is incorrect."
		return 1
	fi
	
	local local_image local_vmlinux local_rootfs local_output
	local_image=$(readlink -f "$1")
	local_vmlinux=$2
	local_rootfs=$3
	local_output=$(readlink -f "$4")
	
	if [ ! -f "${local_image}" ]; then
		echo "Error: local_image=${local_image} is incorrect."
		return 1
	fi
	
	if [ ! -d "${local_output}" ]; then
		echo "Error: local_output=${local_output} is incorrect."
		return 1
	fi
	
	if [ -f "${local_output}/${local_vmlinux}" ]; then
		echo "Error: The ${local_output}/${local_vmlinux} file already exists, please delete the file first."
		return 1
	fi
	
	if [ -f "${local_output}/${local_rootfs}" ]; then
		echo "Error: The ${local_output}/${local_rootfs} file already exists, please delete the file first."
		return 1
	fi
	
	local local_rootfs_offset local_image_offset local_rootfs_size
	local_rootfs_offset=$(grep -oba hsqs ${local_image} | grep -oP '[0-9]*(?=:hsqs)')
	local_image_offset=$(wc -c ${local_image} | grep -oP '[0-9]*(?= )')
	local_rootfs_size=$(expr ${local_image_offset} - ${local_rootfs_offset})
	echo "Info: local_rootfs_offset=${local_rootfs_offset}, local_image_offset=${local_image_offset}, local_rootfs_size=${local_rootfs_size}"
	
	echo "Info: Extract ${local_output}/${local_vmlinux} from ${local_image}, please wait two minutes ..."
	dd if="${local_image}" of="${local_output}/${local_vmlinux}" bs=1 ibs=1 count="${local_rootfs_offset}"
	echo "Info: Extract ${local_output}/${local_rootfs} from ${local_image}, please wait eight minutes ..."
	dd if="${local_image}" of="${local_output}/${local_rootfs}" bs=1 ibs=1 count="${local_rootfs_size}" skip="${local_rootfs_offset}"
	echo "Info: Output ${local_output}/${local_vmlinux} ${local_output}/${local_rootfs}"
	
	return 0
}

unpack_squashfs()
{
	if [ $# -ne 3 ]; then
		echo "Error: parameter is incorrect."
		return 1
	fi
	
	local local_rootfs local_squashfs local_output
	local_rootfs=$(readlink -f "$1")
	local_squashfs=$2
	local_output=$(readlink -f "$3")
	
	if [ ! -f "${local_rootfs}" ]; then
		echo "Error: local_rootfs=${local_rootfs} is incorrect."
		return 1
	fi
	
	if [ ! -d "${local_output}" ]; then
		echo "Error: local_output=${local_output} is incorrect."
		return 1
	fi
	
	local local_unsquashfs='/home/lsgx/develop/lede/builder/staging_dir/host/bin/unsquashfs4'
	if [ ! -f "${local_unsquashfs}" ]; then
		echo "Error: local_unsquashfs=${local_unsquashfs} file does not exist."
		return 1
	fi
	
	if [ -d "${local_output}/${local_squashfs}" ]; then
		echo "Error: The ${local_output}/${local_squashfs} folder directory exists, please delete the directory first."
		return 1
	fi
	
	echo "Info: Unpack ${local_rootfs} to ${local_output}/${local_squashfs}, please wait two minutes ..."
	${local_unsquashfs} -d ${local_output}/${local_squashfs} ${local_rootfs}
	
	return 0
}

pack_squashfs()
{
	if [ $# -ne 3 ]; then
		echo "Error: parameter is incorrect."
		return 1
	fi
	
	local local_squashfs local_rootfs local_output
	local_squashfs=$(readlink -f "$1")
	local_rootfs=$2
	local_output=$(readlink -f "$3")
	
	if [ ! -d "${local_squashfs}" ]; then
		echo "Error: local_squashfs=${local_squashfs} is incorrect."
		return 1
	fi
	
	if [ ! -d "${local_output}" ]; then
		echo "Error: local_output=${local_output} is incorrect."
		return 1
	fi
	
	local local_mksquashfs='/home/lsgx/develop/lede/builder/staging_dir/host/bin/mksquashfs4'
	if [ ! -f "${local_mksquashfs}" ]; then
		echo "Error: local_mksquashfs=${local_mksquashfs} file does not exist."
		return 1
	fi
	
	if [ -f "${local_output}/${local_rootfs}" ]; then
		echo "Error: The ${local_output}/${local_rootfs} file already exists, please delete the file first."
		return 1
	fi
	
	echo "Info: Pack ${local_squashfs} to ${local_output}/${local_rootfs}, please wait two minutes ..."
	${local_mksquashfs} ${local_squashfs} ${local_output}/${local_rootfs} -nopad -noappend -root-owned -comp xz -Xpreset 9 -Xe -Xlc 0 -Xlp 2 -Xpb 2 -b 256k -processors 1
	
	return 0
}

create_image()
{
	if [ $# -ne 4 ]; then
		echo "Error: parameter is incorrect."
		return 1
	fi
	
	local local_vmlinux local_rootfs local_image local_output
	local_vmlinux=$(readlink -f "$1")
	local_rootfs=$(readlink -f "$2")
	local_image=$3
	local_output=$(readlink -f "$4")
	
	if [ ! -f "${local_vmlinux}" ]; then
		echo "Error: local_vmlinux=${local_vmlinux} is incorrect."
		return 1
	fi
	
	if [ ! -f "${local_rootfs}" ]; then
		echo "Error: local_rootfs=${local_rootfs} is incorrect."
		return 1
	fi
	
	if [ ! -d "${local_output}" ]; then
		echo "Error: local_output=${local_output} is incorrect."
		return 1
	fi
	
	local local_padjffs='/home/lsgx/develop/lede/builder/staging_dir/host/bin/padjffs2'
	if [ ! -f "${local_padjffs}" ]; then
		echo "Error: local_padjffs=${local_padjffs} file does not exist."
		return 1
	fi
	
	if [ -f "${local_output}/${local_image}" ]; then
		echo "Error: The ${local_output}/${local_image} file already exists, please delete the file first."
		return 1
	fi
	
	cat ${local_vmlinux} ${local_rootfs} > ${local_output}/${local_image}
	echo "Info: Output ${local_output}/${local_image}"
	${local_padjffs} ${local_output}/${local_image}
	echo "Info: Padding ${local_output}/${local_image}"
	
	return 0
}

if [ "${work_mode}" = "extract" -a "${input_image}" != "" -a "${output_vmlinux}" != "" -a "${output_rootfs}" != "" -a "${output_path}" != "" ]; then
	echo "Extract rootfs from image ..."
	echo "Options: work_mode=${work_mode}, input_image=${input_image}, output_vmlinux=${output_vmlinux}, output_rootfs=${output_rootfs}, output_path=${output_path}"
	extract_rootfs "${input_image}" "${output_vmlinux}" "${output_rootfs}" "${output_path}" || die "Error: Extract rootfs from image failed."
fi

if [ "${work_mode}" = "unpack" -a "${input_rootfs}" != "" -a "${output_squashfs}" != "" -a "${output_path}" != "" ]; then
	echo "Unpack squashfs ..."
	echo "Options: work_mode=${work_mode}, input_rootfs=${input_rootfs}, output_squashfs=${output_squashfs}, output_path=${output_path}"
	unpack_squashfs "${input_rootfs}" "${output_squashfs}" "${output_path}" || die "Error: Unpack squashfs failed."
fi

if [ "${work_mode}" = "pack" -a "${input_squashfs}" != "" -a "${output_rootfs}" != "" -a "${output_path}" != "" ]; then
	echo "Pack squashfs ..."
	echo "Options: work_mode=${work_mode}, input_squashfs=${input_squashfs}, output_rootfs=${output_rootfs}, output_path=${output_path}"
	pack_squashfs "${input_squashfs}" "${output_rootfs}" "${output_path}" || die "Error: Pack squashfs failed."
fi

if [ "${work_mode}" = "create" -a "${input_vmlinux}" != "" -a "${input_rootfs}" != "" -a "${output_image}" != "" -a "${output_path}" != "" ]; then
	echo "Create new image from rootfs and vminux ..."
	echo "Options: work_mode=${work_mode}, input_vmlinux=${input_vmlinux}, input_rootfs=${input_rootfs}, output_image=${output_image}, output_path=${output_path}"
	create_image "${input_vmlinux}" "${input_rootfs}" "${output_image}" "${output_path}" || die "Error: Create new image from rootfs and vminux failed."
fi

exit 0

